import os,sys
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 6)))
sys.path.insert(0, parent_path)

from typing import Any



class OCRDetector(object):
    def __init__(self, keys) -> None:
        from paddleocr import PaddleOCR
        self.keys = keys
        self.detector = PaddleOCR(use_angle_cls = True, lang="ch")

    def __call__(self, data) -> Any:
        input_data = data[self.keys["in"]]
        if len(input_data.shape) == 4:
            input_data = input_data[0]
        result = self.detector.ocr(input_data,cls = True)
        data[self.keys["out"]] = result
        return data

if __name__ == "__main__":

    import cv2, time, glob,pickle
    from core.utils.PIL_draw import pic_text
    def video_test(video_path, save_folder_path):
        height = 1080
        width = 1920
        detector = OCRDetector(
            keys={
                "in":"input_data",
                "out":"ocr",
                }
        )
        cap = cv2.VideoCapture(video_path)
        video_write = cv2.VideoWriter(save_folder_path+"/vis_"+os.path.basename(video_path),cv2.VideoWriter_fourcc(*'XVID'),25,(width,height))
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame = cv2.resize(frame,[width,height])
            data = detector({
                "input_data": frame
            })
            det = data["ocr"]
            for ocr_data in det:
                if ocr_data is None:
                    continue
                for ocr_ in ocr_data:
                    point = [int(x) for x in ocr_[0][0]]
                    frame = pic_text(frame,ocr_[1][0],point=point,fillColor=(0,0,255),size=30)
            video_write.write(frame)
        video_write.release()
            

    def pic_test(file_path):
        detector = OCRDetector(
            keys={
                "in":"input_data",
                "out":"ocr",
                }
        )
        if os.path.isdir(file_path):
            file_path = glob.glob(file_path+"/*")
        else:
            file_path = [file_path]
        for path_ in file_path:
            frame = cv2.imread(path_)
            # frame = cv2.resize(frame,[1920,1080])
            file_name = os.path.basename(path_)
            data = detector({
                "input_data": frame
            })
            det = data["ocr"]
            with open("temp.pkl","wb") as f:
                pickle.dump(det,f)
            for ocr_data in det:
                if ocr_data is None:
                    continue
                for ocr_ in ocr_data:
                    point = [int(x) for x in ocr_[0][0]]
                    frame = pic_text(frame,ocr_[1][0],point=point,fillColor=(0,0,255))
            cv2.imwrite("data/output/{}".format(file_name),frame)
    pic_test("data/input/ocr")
    # video_test("data/input/baixiang/2.mp4",
            #    "data/output/baixiang")